#include <iostream>
#include <vector>
#include <map>
#include <set>
#include <algorithm>

#include <cmath>

#include <matrix.h>
#include <cluster.h>

#include <sds/sfile.h>
#include <sds/sdataset.h>
#include <sds/sdataview.h>

#include <lemon/random.h>

using namespace std;
using namespace sds;

#if !defined(INIT_NORMAL) && !defined(INIT_ADV)
#define INIT_ADV
#endif


// parameters

int smallClusterSize = 230; //256
double smallDistance = 1.44; //1.64

double maxDistance = 128.0;
double maxDistanceLast = 96.0;

int stepNum = 10;

int initialCluster = 60;

// global variables


double invMaxDistance = inv_map_euclid(maxDistance);
double invMaxDistanceLast = inv_map_euclid(maxDistanceLast);

class IntMap : public std::vector<int> {
public:
  typedef int Key;

  IntMap(int size, int value) : std::vector<int>(size, value) {}
  void set(int key, int value) { (*this)[key] = value;}
};

void classify(const SDataView<Point>& src,
	      std::vector<double>& probs,
	      std::vector<Cluster>& clusters) {

  vector<Vector> expected;
  vector<Matrix> inverse;
  vector<double> constcomp;

  for (int k = 0; k < clusters.size(); ++k) {
    Vector exp(CH);
    for (int i = 0; i < CH; ++i) {
      exp(i) = clusters[k].expected(i);
    }
    expected.push_back(exp);
    Matrix cov(CH, CH);
    for (int i = 0; i < CH; ++i) {
      for (int j = 0; j < CH; ++j) {
	cov(i, j) = clusters[k].covariance(i, j);
      }
    }
    inverse.push_back(cov.inverse());
    constcomp.push_back(log(probs[k]) - 0.5 * log(fabs(cov.determinant())) -
			0.5 * CH * log(2 * M_PI));
  }

  for (int k = 0; k < clusters.size(); ++k) {
    clusters[k].clear();
  }

  SDataView<Point>::Data::const_iterator it;
  for (it = src.data().begin(); it != src.data().end(); ++it) {
    if (it->zero()) continue;
    Vector current(CH);
    for (int l = 0; l < CH; ++l) {
      current(l) = (*it)[l];
    }
    for (int k = 0; k < clusters.size(); ++k) {
      Vector diff = current - expected[k];
      double dist = - 0.5 * diff * inverse[k] * diff + constcomp[k];
     
      clusters[k].add(*it, exp(dist));	
    }
  }
  
  double sum = 0.0;
  for (int k = 0; k < clusters.size(); ++k) {
    sum += clusters[k].size();
  }
  for (int k = 0; k < clusters.size(); ++k) {
    probs[k] = clusters[k].size() / sum;
  }
}

void mindistCluster(const SDataView<Point>& src,
		    const std::vector<Point>& centers,
		    std::vector<double>& probs,
		    std::vector<Cluster>& clusters) {
  SDataView<Point>::Data::const_iterator it;
  clusters.resize(centers.size());
  probs.resize(centers.size());
  for (it = src.data().begin(); it != src.data().end(); ++it) {
    if (it->zero()) continue;
    int minPos = 0;
    double minVal = mon_distance_euclid(*it, centers[0]);
    for (int k = 1; k < (int)centers.size(); ++k) {
      double val = mon_distance_euclid(*it, centers[k]);
      if (val < minVal) {
	minVal = val;
	minPos = k;
      }
    }
    clusters[minPos] += *it;
  }

  double sum = 0.0;
  for (int k = 0; k < centers.size(); ++k) {
    sum += clusters[k].size();
  }

  for (int k = 0; k < centers.size(); ++k) {
    probs[k] += clusters[k].size() / sum;
  }
  
}

#if defined(INIT_NORMAL)
void start(const SDataView<Point>& src,
	   std::vector<double>& probs,
	   std::vector<Cluster>& clusters) {
  Cluster info;
  SDataView<Point>::Data::const_iterator it;
  for (it = src.data().begin(); it != src.data().end(); ++it) {
    if (it->zero()) continue;
    info += *it;
  }
  Vector expected(CH);
  Matrix covariance(CH, CH);
  for (int i = 0; i < CH; ++i) {
    for (int j = 0; j < CH; ++j) {
      covariance(i, j) = info.covariance(i, j);
    }
    expected(i) = info.expected(i);
  }
  std::vector<Point> centers;  

  Matrix cholesky = covariance.cholesky();
  for (int k = 0; k < initialCluster; ++k) {
    Vector v(CH);
    for (int i = 0; i < CH; ++i) {
      v(i) = lemon::rnd.gauss();
    }
    //    v = exp((double)k / (double)initialCluster) * (v * cholesky) + expected;
    v = v * cholesky + expected;
    Point p;
    for (int i = 0; i < CH; ++i) {
      p[i] = v(i);
    }
    centers.push_back(p);
  }

  mindistCluster(src, centers, probs, cluster);
}

#elif defined(INIT_ADV)

void start(const SDataView<Point>& src, 
	   std::vector<double>& probs,
	   std::vector<Cluster>& clusters) {

  double dss;
  SDataView<double> ds(src.getSize());

  std::vector<Point> centers;
  {
    int n = lemon::rnd[src.data().size()];
    while (src.data()[n].zero()) {
      n = lemon::rnd[src.data().size()];
    }
    Point cntr = src.data()[n];
    centers.push_back(cntr);

    dss = 0.0;
    for (int i = 0; i < src.data().size(); ++i) {
      if (src.data()[i].zero()) continue;
      ds.data()[i] = mon_distance_euclid(src.data()[i], cntr);
      dss += ds.data()[i];
    }
  }

  for (int k = 1; k < initialCluster; ++k) {
    std::cerr << k << std::endl;
    double r = lemon::rnd(dss);

    int n;
    for (n = 0; ; ++n) {
      if (src.data()[n].zero()) continue;
      if (r < ds.data()[n]) break;
      r -= ds.data()[n];
    }

    Point cntr = src.data()[n];
    centers.push_back(cntr);

    dss = 0.0;
    for (int i = 0; i < src.data().size(); ++i) {
      if (src.data()[i].zero()) continue;
      ds.data()[i] = 
	std::min(ds.data()[i], mon_distance_euclid(src.data()[i], cntr));
      dss += ds.data()[i];
    }
  }
  mindistCluster(src, centers, probs, clusters);
}

#endif

void Read(std::string filename, SDataView<Point>& src, vector<SInteger>& size) {
  SFile file(filename);

  SDataView<unsigned char> views[CH];
  
  for (int i = 0; i < CH; ++i) { 
    SDataset dataset = file.getDataset(i);
    views[i] = dataset.getDataView<unsigned char>();
  }

  size = views[0].getSize();

  src = SDataView<Point>(size);

  for (int i = 0; i < size[0]; ++i) {
    for (int j = 0; j < size[1]; ++j) {
      Point item;
      for (int k = 0; k < CH; ++k) {
	item[k] = (double)views[k][i][j];
      }
      src[i][j] = item;
    }
  } 
}

void Write(std::string targetName, 
	   const SDataView<Point>& src,
	   const std::vector<double>& probs, 
	   const std::vector<Cluster>& clusters, 
	   const vector<SInteger>& size) {
  SFile target(targetName, SFile::CREATE);
  SDataset targetDataset = target.createDataset("clustered", UINT8, size);

  SDataView<unsigned char> targetView(size);

  vector<Vector> expected;
  vector<Matrix> inverse;
  vector<double> constcomp;

  for (int k = 0; k < clusters.size(); ++k) {
    Vector exp(CH);
    for (int i = 0; i < CH; ++i) {
      exp(i) = clusters[k].expected(i);
    }
    expected.push_back(exp);
    Matrix cov(CH, CH);
    for (int i = 0; i < CH; ++i) {
      for (int j = 0; j < CH; ++j) {
	cov(i, j) = clusters[k].covariance(i, j);
      }
    }
    inverse.push_back(cov.inverse());
    constcomp.push_back(log(probs[k]) - 0.5 * log(fabs(cov.determinant())) -
			0.5 * CH * log(2 * M_PI));
  }

  for (int i = 0; i < size[0]; ++i) {
    for (int j = 0; j < size[1]; ++j) {
      targetView[i][j] = 0;
      if (((const Point&)src[i][j]).zero()) continue;

      Vector current(CH);
      for (int l = 0; l < CH; ++l) {
	current(l) = ((const Point&)src[i][j])[l];
      }

      int maxPos = -1; 
      double maxVal = numeric_limits<double>::max();
      for (int k = 0; k < clusters.size(); ++k) {
	Vector diff = current - expected[k];
	double dist = - 0.5 * diff * inverse[k] * diff + constcomp[k];
	
	if (dist < maxVal) {
	  maxPos = k;
	  maxVal = dist;
	}
      }
      targetView[i][j] = (unsigned char)maxPos + 1;
    }
  }
  targetDataset.setDataView<unsigned char>(targetView);  
}

int main(int argc, const char *argv[]) {
  lemon::rnd.seed(time(0));

  assert(argc == 3);
  SDataView<Point> src;
  vector<SInteger> size;
  Read(argv[1], src, size);
  std::vector<double> probs;
  std::vector<Cluster> clusters;
  start(src, probs, clusters);
  for (int step = 0; step < stepNum; ++step) {
    std::cout << (double)step / stepNum << std::endl;
    classify(src, probs, clusters);
  }
  Write(argv[2], src, probs, clusters, size);

  return 0;
}
